/******************************************************************************
 * Copyright 2022 The Untouch Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *****************************************************************************/

#include "modules/calibration/lidar_imu_calibrator/include/lidar_imu_calibrator.h"

#include <QFileDialog>
#include <QMessageBox>
#include <fstream>

#include <glk/lines.hpp>
#include <glk/pointcloud_buffer.hpp>
#include <glk/primitives/primitives.hpp>

#include "yaml-cpp/yaml.h"

#include "cyber/common/log.h"
#include "modules/calibration/lidar_imu_calibrator/include/compensate_data_dialog.h"
#include "modules/calibration/lidar_imu_calibrator/include/load_data_dialog.h"
#include "modules/calibration/lidar_imu_calibrator/include/record_data_dialog.h"
#include "modules/calibration/lidar_imu_calibrator/include/save_extrinsic_parameters_dialog.h"
#include "modules/localization/msf/common/io/velodyne_utility.h"

namespace lidar_imu_calibrator {

const double ARC_TO_DEG = 57.29577951308238;
const double DEG_TO_ARC = 0.0174532925199433;

LidarImuCalibratorApplication::LidarImuCalibratorApplication()
    : Application() {}
LidarImuCalibratorApplication::~LidarImuCalibratorApplication() {}

bool LidarImuCalibratorApplication::init(const char* window_name,
                                         const Eigen::Vector2i& size,
                                         const char* glsl_version) {
  if (!Application::init(window_name, size, glsl_version)) {
    return false;
  }

  framebuffer_size_ = size;
  progress_.reset(new guik::ProgressModal("progress modal"));

  std::string data_directory =
      "/apollo/modules/calibration/lidar_imu_calibrator/data";

  main_canvas_.reset(new guik::GLCanvas(data_directory, size));
  if (!main_canvas_->ready()) {
    close();
  }

  return true;
}

void LidarImuCalibratorApplication::draw_ui() {
  main_canvas_->draw_ui();
  DrawKeyFrameSettingDialog(key_frame_parameter_[0], odometry_frame_set_[0]);
  DrawKeyFrameSettingDialog(key_frame_parameter_[1], odometry_frame_set_[1]);
  MainMenu();
  MouseControl();
}

void LidarImuCalibratorApplication::draw_gl() {
  glEnable(GL_VERTEX_PROGRAM_POINT_SIZE);

  main_canvas_->bind();
  main_canvas_->shader->set_uniform("color_mode", 2);
  main_canvas_->shader->set_uniform(
      "model_matrix",
      (Eigen::UniformScaling<float>(3.0f) * Eigen::Isometry3f::Identity())
          .matrix());

  const auto& coord = glk::Primitives::instance()->primitive(
      glk::Primitives::COORDINATE_SYSTEM);
  coord.draw(*main_canvas_->shader);

  main_canvas_->shader->set_uniform("color_mode", 1);
  main_canvas_->shader->set_uniform(
      "model_matrix", (Eigen::Translation3f(Eigen::Vector3f::UnitZ() * -0.02) *
                       Eigen::Isometry3f::Identity())
                          .matrix());
  main_canvas_->shader->set_uniform("material_color",
                                    Eigen::Vector4f(0.8f, 0.8f, 0.8f, 1.0f));
  const auto& grid =
      glk::Primitives::instance()->primitive(glk::Primitives::GRID);
  grid.draw(*main_canvas_->shader);

  main_canvas_->shader->set_uniform("point_scale", 1.0f);
  for (auto& odometry_frame_set : odometry_frame_set_) {
    if (odometry_frame_set) {
      odometry_frame_set->DrawPointCloud(*main_canvas_->shader);
    }
  }

  main_canvas_->unbind();
  main_canvas_->render_to_screen();

  glDisable(GL_VERTEX_PROGRAM_POINT_SIZE);
}

void LidarImuCalibratorApplication::framebuffer_size_callback(
    const Eigen::Vector2i& size) {
  main_canvas_->set_size(size);
  framebuffer_size_ = size;
}

void LidarImuCalibratorApplication::ResetKeyFrameParameter(
    const Eigen::Isometry3d& extrinsic,
    std::shared_ptr<KeyFrameParameter>& key_frame_parameter,
    const std::string& title) {
  if (!key_frame_parameter) {
    key_frame_parameter = std::make_shared<KeyFrameParameter>();
  }
  key_frame_parameter->title_ = title;
  key_frame_parameter->delay_time_ = 0.0f;
  key_frame_parameter->delta_updated_ = false;
  key_frame_parameter->key_frame_delta_x_ = 0.1f;
  key_frame_parameter->key_frame_delta_angle_ = 1.0f;
  key_frame_parameter->x_ = extrinsic.translation().x();
  key_frame_parameter->y_ = extrinsic.translation().y();
  key_frame_parameter->z_ = extrinsic.translation().z();
  Eigen::Quaterniond qua(extrinsic.linear());
  auto euler = qua.toRotationMatrix().eulerAngles(2, 1, 0);
  key_frame_parameter->yaw_ = euler[0] * ARC_TO_DEG;
  key_frame_parameter->pitch_ = euler[1] * ARC_TO_DEG;
  key_frame_parameter->roll_ = euler[2] * ARC_TO_DEG;
}

void LidarImuCalibratorApplication::DrawKeyFrameSettingDialog(
    std::shared_ptr<KeyFrameParameter>& key_frame_parameter,
    std::shared_ptr<OdometryFrameSet>& odometry_frame_set) {
  if (!key_frame_parameter) {
    return;
  }
  ImGui::Begin(key_frame_parameter->title_.c_str(), nullptr,
               ImGuiWindowFlags_AlwaysAutoResize);

  bool updated = false;
  updated |= ImGui::DragFloat(
      "delta x", &key_frame_parameter->key_frame_delta_x_, 0.1f, 0.1f, 100.0f);
  updated |= ImGui::DragFloat("delta angle",
                              &key_frame_parameter->key_frame_delta_angle_,
                              0.01f, 0.01f, 3.15f);
  updated |=
      ImGui::DragFloat("x", &key_frame_parameter->x_, 0.01f, 0.01f, 100.0f);
  updated |=
      ImGui::DragFloat("y", &key_frame_parameter->y_, 0.01f, 0.01f, 100.0f);
  updated |=
      ImGui::DragFloat("z", &key_frame_parameter->z_, 0.01f, 0.01f, 100.0f);
  updated |= ImGui::DragFloat("roll", &key_frame_parameter->roll_, 0.001f,
                              0.001f, 3.15f);
  updated |= ImGui::DragFloat("pitch", &key_frame_parameter->pitch_, 0.001f,
                              0.001f, 3.15f);
  updated |= ImGui::DragFloat("yaw", &key_frame_parameter->yaw_, 0.001f, 0.001f,
                              3.15f);
  updated |= ImGui::DragFloat("delay time", &key_frame_parameter->delay_time_,
                              0.0001f, 0.0001f, 1.0f);

  if (key_frame_parameter->delta_updated_ && !updated) {
    Eigen::Isometry3d extrinsics = Eigen::Isometry3d::Identity();
    extrinsics.translation() =
        Eigen::Vector3d(key_frame_parameter->x_, key_frame_parameter->y_,
                        key_frame_parameter->z_);
    extrinsics.linear() =
        Eigen::AngleAxisd(key_frame_parameter->yaw_ * DEG_TO_ARC,
                          Eigen::Vector3d::UnitZ()) *
        Eigen::AngleAxisd(key_frame_parameter->pitch_ * DEG_TO_ARC,
                          Eigen::Vector3d::UnitY()) *
        Eigen::AngleAxisd(key_frame_parameter->roll_ * DEG_TO_ARC,
                          Eigen::Vector3d::UnitX())
            .toRotationMatrix();
    odometry_frame_set->SelectKeyFrames(
        key_frame_parameter->key_frame_delta_x_,
        key_frame_parameter->key_frame_delta_angle_);
    odometry_frame_set->UpdateDelayTime(key_frame_parameter->delay_time_);
    odometry_frame_set->UpdateExtrinsics(extrinsics);
  }
  key_frame_parameter->delta_updated_ = updated;
  ImGui::End();
}

void LidarImuCalibratorApplication::MainMenu() {
  ImGui::BeginMainMenuBar();

  bool record_data = false;
  bool load_first_data = false;
  bool load_second_data = false;
  bool compensate_first_data = false;
  bool compensate_second_data = false;
  bool save_dialog = false;
  if (ImGui::BeginMenu("Operate")) {
    if (ImGui::MenuItem("Record Lidar IMU Data")) {
      record_data = true;
    }

    if (ImGui::BeginMenu("First Lidar IMU")) {
      if (ImGui::MenuItem("Load Data")) {
        load_first_data = true;
      }

      if (ImGui::MenuItem("Compensate Data")) {
        compensate_first_data = true;
      }

      if (ImGui::MenuItem("Calibrate")) {
        std::cout << "star_cali" << std::endl;
        SelfCalibration(0);
      }

      ImGui::EndMenu();
    }

    if (ImGui::BeginMenu("Second Lidar IMU")) {
      if (ImGui::MenuItem("Load Data")) {
        load_second_data = true;
      }

      if (ImGui::MenuItem("Compensate Data")) {
        compensate_second_data = true;
      }
      ImGui::EndMenu();
    }

    if (ImGui::MenuItem("Save Extrinsic Parameter")) {
      save_dialog = true;
    }

    if (ImGui::MenuItem("Quit")) {
      close();
    }

    ImGui::EndMenu();
  }

  RecordData(record_data);
  LoadFirstLidarImuData(load_first_data);
  CompensateFirstLidarImuData(compensate_first_data);
  LoadSecondLidarImuData(load_second_data);
  CompensateSecondLidarImuData(compensate_second_data);
  Save(save_dialog);

  if (ImGui::BeginMenu("View")) {
    if (ImGui::MenuItem("Reset Camera")) {
      main_canvas_->reset_camera();
    }
    if (ImGui::MenuItem("Projection Setting")) {
      main_canvas_->show_projection_setting();
    }
    ImGui::EndMenu();
  }

  bool show_version = false;
  if (ImGui::BeginMenu("Help")) {
    if (ImGui::MenuItem("About")) {
      show_version = true;
    }

    ImGui::EndMenu();
  }

  if (show_version) {
    version_modal_->Open();
  }
  version_modal_->Run();

  ImGui::EndMainMenuBar();
}

void LidarImuCalibratorApplication::RecordData(bool record_data) {
  if (!record_data) {
    return;
  }

  RecordDataDialog dialog;
  dialog.exec();
}

void LidarImuCalibratorApplication::LoadFirstLidarImuData(bool load) {
  if (!load) {
    return;
  }

  std::shared_ptr<OdometryFrameSet> odometry_frame_set(new OdometryFrameSet());
  LoadDataDialog dialog(odometry_frame_set);
  if (dialog.exec() == QDialog::Accepted) {
    odometry_frame_set_[0] = odometry_frame_set;
    ResetKeyFrameParameter(odometry_frame_set_[0]->GetExtrinsicParameter(),
                           key_frame_parameter_[0], "first key frame settings");
    odometry_frame_set_[0]->SelectKeyFrames(
        key_frame_parameter_[0]->key_frame_delta_x_,
        key_frame_parameter_[0]->key_frame_delta_angle_);
    key_frame_parameter_[0]->delta_updated_ = true;
  }
}

void LidarImuCalibratorApplication::CompensateFirstLidarImuData(
    bool compensate) {
  if (!compensate) {
    return;
  }

  if (!odometry_frame_set_[0] || !odometry_frame_set_[0]->IsNeedCompensate()) {
    QMessageBox::warning(
        nullptr, QStringLiteral("Prompt"),
        QStringLiteral("No point data need to be compensated!"));
    return;
  }

  CompensateDataDialog dialog(odometry_frame_set_[0]);
  if (dialog.exec() == QDialog::Accepted) {
    odometry_frame_set_[0]->SelectKeyFrames(
        key_frame_parameter_[0]->key_frame_delta_x_,
        key_frame_parameter_[0]->key_frame_delta_angle_);
    key_frame_parameter_[0]->delta_updated_ = true;
  }
}

void LidarImuCalibratorApplication::LoadSecondLidarImuData(bool load) {
  if (!load) {
    return;
  }

  std::shared_ptr<OdometryFrameSet> odometry_frame_set(new OdometryFrameSet());
  LoadDataDialog dialog(odometry_frame_set);
  if (dialog.exec() == QDialog::Accepted) {
    odometry_frame_set_[1] = odometry_frame_set;
    ResetKeyFrameParameter(odometry_frame_set_[1]->GetExtrinsicParameter(),
                           key_frame_parameter_[1],
                           "second key frame settings");
    odometry_frame_set_[1]->SelectKeyFrames(
        key_frame_parameter_[1]->key_frame_delta_x_,
        key_frame_parameter_[1]->key_frame_delta_angle_);
    key_frame_parameter_[1]->delta_updated_ = true;
  }
}

void LidarImuCalibratorApplication::CompensateSecondLidarImuData(
    bool compensate) {
  if (!compensate) {
    return;
  }

  if (!odometry_frame_set_[1] || !odometry_frame_set_[1]->IsNeedCompensate()) {
    QMessageBox::warning(
        nullptr, QStringLiteral("Prompt"),
        QStringLiteral("No point data need to be compensated!"));
    return;
  }

  CompensateDataDialog dialog(odometry_frame_set_[1]);
  if (dialog.exec() == QDialog::Accepted) {
    odometry_frame_set_[1]->SelectKeyFrames(
        key_frame_parameter_[1]->key_frame_delta_x_,
        key_frame_parameter_[1]->key_frame_delta_angle_);
    key_frame_parameter_[1]->delta_updated_ = true;
  }
}

void LidarImuCalibratorApplication::Save(bool save_dialog) {
  if (!save_dialog) {
    return;
  }

  QVector<int> extrinsic_paramter_Type_indexs;
  if (key_frame_parameter_[0]) {
    extrinsic_paramter_Type_indexs.push_back(0);
  }

  if (key_frame_parameter_[1]) {
    extrinsic_paramter_Type_indexs.push_back(1);
  }

  SaveExtrinsicParametersDialog dialog(extrinsic_paramter_Type_indexs);
  if (dialog.exec() == QDialog::Accepted) {
    if (SaveExtrinsicParameter(dialog.GetSavePath().toStdString(),
                               dialog.GetExtrinsicParameterTypeIndex())) {
      QMessageBox::warning(nullptr, QStringLiteral("Prompt"),
                           QStringLiteral("Save extrinsic parameter success!"));
    } else {
      QMessageBox::warning(nullptr, QStringLiteral("Prompt"),
                           QStringLiteral("Save extrinsic parameter fail!"));
    }
  }
}

void LidarImuCalibratorApplication::MouseControl() {
  ImGuiIO& io = ImGui::GetIO();
  if (!io.WantCaptureMouse) {
    main_canvas_->mouse_control();
  }
}

bool LidarImuCalibratorApplication::SaveExtrinsicParameter(
    std::string file_path, int index) {
  Eigen::Isometry3d extrinsics = Eigen::Isometry3d::Identity();
  extrinsics.translation() = Eigen::Vector3d(key_frame_parameter_[index]->x_,
                                             key_frame_parameter_[index]->y_,
                                             key_frame_parameter_[index]->z_);
  extrinsics.linear() =
      Eigen::AngleAxisd(key_frame_parameter_[index]->yaw_ * DEG_TO_ARC,
                        Eigen::Vector3d::UnitZ()) *
      Eigen::AngleAxisd(key_frame_parameter_[index]->pitch_ * DEG_TO_ARC,
                        Eigen::Vector3d::UnitX()) *
      Eigen::AngleAxisd(key_frame_parameter_[index]->roll_ * DEG_TO_ARC,
                        Eigen::Vector3d::UnitY())
          .toRotationMatrix();

  return odometry_frame_set_[index]->SaveExtrinsicParameter(extrinsics,
                                                            file_path);
}

bool LidarImuCalibratorApplication::SelfCalibration(int index_odom){
  int index = 1 == index_odom ? 1 : 0;

  std::vector<OdometryFrame::Ptr> frames;
  odometry_frame_set_[index]->GetFrames(frames);
  double error_pre = 0, error;
  int loop_times = 0;
  while (true)
  {
    Eigen::Isometry3d extrinsics = odometry_frame_set_[index]->GetExtrinsicParameter();
    if (0 == error_pre || fabs((error_pre - error) / (error_pre + 10e-6)) < 0.01) {
      if (loop_times >= 2) break;
      loop_times ++;

      error = self_clibrator_.Calibrate(frames, extrinsics);
      if (-1 == error) {
        error_pre = error;
        continue;
      }
      odometry_frame_set_[index]->UpdateExtrinsics(extrinsics);
      CompensateFirstLidarImuData(true);
      std::string savePath1 = "/apollo/data/lidar_imu/result.pcd";
      self_clibrator_.SaveCloud(savePath1, extrinsics);
    }
  }
      return true;

  return false;
}
}  // namespace lidar_imu_calibrator

